import time

import boto3

athena = boto3.client('athena')


def get_rs_status(eid: str):
    response = athena.get_query_execution(
        QueryExecutionId=eid
    )
    return response['QueryExecution']['Status']['State']


def lambda_handler(event, context):
    model_number: str = event['model']
    model_number = model_number.replace(',', '|')
    sql = f" select material from dim.material where regexp_extract(model,'({model_number})') is not null and type_code = 'FRIDGE' "

    start_rs = athena.start_query_execution(
        QueryString=sql,
        WorkGroup='application'
    )

    execution_id = start_rs['QueryExecutionId']
    state = get_rs_status(execution_id)
    while state == 'QUEUED' or state == 'RUNNING':
        time.sleep(2)
        state = get_rs_status(execution_id)

    if state == 'FAILED' or state == 'CANCELLED':
        return {
            'statusCode': 500
        }

    rs = athena.get_query_results(
        QueryExecutionId=execution_id,
        MaxResults=1000
    )

    rows = rs['ResultSet']['Rows']
    materials = []
    for row in rows:
        if row['Data']['VarCharValue'] == 'material':
            continue
        else:
            materials.append(row['Data'][0]['material'])

    ms = ','.join(materials)

    return {
        'statusCode': 200,
        'materials': ms
    }
